import pygame
import json
import math

def draw_path_on_map(map_image_path, output_path):
    # 初始化pygame
    pygame.init()

    # 加载地图图片
    map_image = pygame.image.load(map_image_path)
    original_width, original_height = map_image.get_size()

    # 设置窗口为全屏
    screen = pygame.display.set_mode((0, 0), pygame.FULLSCREEN)
    pygame.display.set_caption("Map Path Drawer")
    screen_width, screen_height = pygame.display.get_surface().get_size()

    # 计算地图显示尺寸 - 保持原始比例，贴着左边缘
    # 高度不超过屏幕高度
    scale_factor = min(1.0, screen_height / original_height)
    display_width = int(original_width * scale_factor)
    display_height = int(original_height * scale_factor)

    # 缩放地图图片
    scaled_map_image = pygame.transform.scale(map_image, (display_width, display_height))

    # 地图位置 - 贴着左边缘
    map_x = 0
    map_y = (screen_height - display_height) // 2  # 垂直居中

    # 清屏并显示地图
    screen.fill((0, 0, 0))  # 填充黑色背景
    screen.blit(scaled_map_image, (map_x, map_y))
    pygame.display.flip()

    # 路径点列表
    path_points = []  # 存储路径点的原始坐标
    scaled_path_points = []  # 存储路径点的缩放后坐标（仅用于绘制）
    map_change_points = {}  # 存储更换地图的指定位置（原始坐标）及对应地图索引

    # 主循环
    running = True
    drawing = True
    current_map_index = 0  # 当前地图索引
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN and drawing:
                # 获取鼠标点击的位置
                mouse_x, mouse_y = event.pos

                # 检查点击是否在地图区域内
                if map_x <= mouse_x < map_x + display_width and map_y <= mouse_y < map_y + display_height:
                    # 将鼠标坐标映射回原始地图分辨率
                    original_x = int((mouse_x - map_x) / scale_factor)
                    original_y = int((mouse_y - map_y) / scale_factor)

                    path_points.append((original_x, original_y))
                    scaled_path_points.append((mouse_x, mouse_y))  # 存储缩放后的坐标（仅用于绘制）

                    # 在窗口中绘制点
                    pygame.draw.circle(screen, (255, 0, 0), (mouse_x, mouse_y), 5)
                    # 如果有多于一个点，绘制线段
                    if len(path_points) > 1:
                        prev_x, prev_y = scaled_path_points[-2]
                        pygame.draw.line(screen, (255, 0, 0), (prev_x, prev_y), (mouse_x, mouse_y), 2)
                    pygame.display.flip()
            elif event.type == pygame.MOUSEBUTTONDOWN and not drawing:
                # 在路径上选择更换地图的指定位置
                mouse_x, mouse_y = event.pos
                for i, (scaled_x, scaled_y) in enumerate(scaled_path_points):
                    # 计算鼠标点击位置与路径点的距离
                    distance = math.sqrt((mouse_x - scaled_x) ** 2 + (mouse_y - scaled_y) ** 2)
                    if distance <= 10:  # 如果距离小于等于10像素
                        # 将该路径点添加到更换地图的指定位置
                        map_change_points[path_points[i]] = current_map_index + 1  # 记录坐标和对应地图索引
                        # 在路径点的红点上叠加一个蓝色的点
                        pygame.draw.circle(screen, (0, 0, 255), (scaled_x, scaled_y), 5)
                        pygame.display.flip()
                        current_map_index += 1  # 增加地图索引
                        break
            elif event.type == pygame.KEYDOWN and event.key == pygame.K_SPACE:
                # 按空格键结束绘制
                if drawing:
                    drawing = False
                else:
                    # 保存路径点和更换地图的指定位置（仅原始坐标和图片分辨率）
                    with open(output_path, 'w') as f:
                        # 将字典的键转换为字符串以支持JSON序列化
                        serialized_map_change_points = {str(k): v for k, v in map_change_points.items()}
                        json.dump({
                            "path_points": path_points,
                            "map_change_points": serialized_map_change_points,
                            "image_resolution": [original_width, original_height]  # 仅保存图片原始分辨率
                        }, f)
                    running = False

    # 退出pygame
    pygame.quit()

if __name__ == "__main__":
    map_image_path = "assets/maps/map5.png"
    output_path = "path_points.json"
    draw_path_on_map(map_image_path, output_path)
    print(f"Path points and map data saved to {output_path}")